#include "common.h"
#include <apr_network_io.h>
#include <apr_strings.h>

#define LUA_SOCK_CHUNK_SIZE 8192

/** A socket, its address, and all related metadata
 */
typedef struct lua_apr_socket {
    apr_socket_t *socket;
    apr_sockaddr_t *addr;
    apr_pool_t *pool;
} lua_apr_socket;

/** Create and connect a client TCP socket to a given host/port combination
 * \luaparam host the IP address to connect to
 * \luaparam port the port number to connect to
 * \luareturn the raw lua_apr_socket struct for the socket
 */
static int lua_apr_socket_connect(lua_State *L) {
    const char *host = luaL_checkstring(L, 1);
    apr_port_t port = (apr_port_t) luaL_checkint(L, 2);
    apr_status_t rv;
    apr_pool_t *pool = NULL;
    
    rv = apr_pool_create(&pool, NULL);
    if(rv != APR_SUCCESS) {
        return luaL_error(L, "An error occurred initializing memory");
    }
    
    apr_socket_t *socket = NULL;
    rv = apr_socket_create(&socket, APR_INET, SOCK_STREAM, APR_PROTO_TCP, pool);
    if(rv != APR_SUCCESS) {
        return luaL_error(L, "An error occurred creating a new socket");
    }
    
    apr_sockaddr_t *sa = NULL;
    rv = apr_sockaddr_info_get(&sa, host, APR_UNSPEC, port, APR_IPV4_ADDR_OK, pool);
    if(rv != APR_SUCCESS) {
        return luaL_error(L, "An error occurred getting the socket info");
    }
    
    rv = apr_socket_connect(socket, sa);
    if(rv != APR_SUCCESS) {
        lua_pushnil(L);
        lua_pushfstring(L, "Connection to %s:%d failed", host, port);
        return 2;
    }
    
    lua_apr_socket *sock = (lua_apr_socket *) lua_newuserdata(L, sizeof(lua_apr_socket));
    sock->socket = socket;
    sock->addr = sa;
    sock->pool = pool;
    
    luaL_getmetatable(L, "apr.net.socket");
    lua_setmetatable(L, -2);
    return 1;
}

/** Close the socket and clean up memory
 * \luaparam sock the socket to close and clean
 */
static int lua_apr_socket_close(lua_State *L) {
    lua_apr_socket *sock = luaL_checkudata(L, 1, "apr.net.socket");
    if(sock->socket == NULL) {
        return 0;
    }
    
    apr_socket_close(sock->socket);
    apr_pool_destroy(sock->pool);
    sock->socket = NULL;
    sock->addr = NULL;
    return 0;
}

/** If given a host name, find all IP addresses that resolve to that host and return them as a table
 * \luaparam host the host name to check - either an IP address or a domain name
 * \luareturn the IP addresses that resolve to this host as a table
 */
static int lua_apr_socket_dnsresolve(lua_State *L) {
    const char *host = luaL_checkstring(L, 1);
    apr_sockaddr_t *sa;
    apr_pool_t *pool = NULL;
    apr_status_t rv;
    
    rv = apr_pool_create(&pool, NULL);
    if(rv != APR_SUCCESS) {
        return luaL_error(L, "An error occurred initializing memory");
    }
    
    rv = apr_sockaddr_info_get(&sa, host, APR_UNSPEC, 0, APR_IPV4_ADDR_OK, pool);
    if(rv != APR_SUCCESS) {
        lua_pushnil(L);
        lua_pushfstring(L, "Invalid IP address or host: %s", host);
        apr_pool_destroy(pool);
        return 2;
    }
    
    char *addr;
    rv = apr_sockaddr_ip_get(&addr, sa);
    if(rv != APR_SUCCESS) {
        lua_pushnil(L);
        lua_pushfstring(L, "Invalid IP address or host: %s", host);
        apr_pool_destroy(pool);
        return 2;
    }

    lua_newtable(L);
    lua_pushstring(L, addr);
    lua_rawseti(L, -2, 1);
    
    apr_pool_destroy(pool);
    return 1;
}

/** Send a request across a client socket
 * \luaparam sock the raw socket struct to send across
 * \luaparam data the data to send
 */
static int lua_apr_socket_send(lua_State *L) {
    lua_apr_socket *sock = (lua_apr_socket *) luaL_checkudata(L, 1, "apr.net.socket");
    apr_status_t rv;
    apr_size_t datalen = 0;
    int argnum = lua_gettop(L);
    const char *data = NULL;
    int argindex = 0;
    
    for(argindex=2; argindex<=argnum; argindex++) {
        data = luaL_checklstring(L, argindex, &datalen);
        rv = apr_socket_send(sock->socket, data, &datalen);
        
        if(rv != APR_SUCCESS) {
            return luaL_error(L, "An error occurred sending to socket");
        }
    }
    
    return 0;
}

/** Receive the next chunk from a client socket, returning the chunk or nil if no chunks are left
 * \luaupval sock the socket to receive from - from the closure
 * \luareturn the next chunk
 */
static int lua_apr_socket_recv(lua_State *L) {
    lua_apr_socket *sock = (lua_apr_socket *) luaL_checkudata(L, 1, "apr.net.socket");
    apr_status_t rv = APR_SUCCESS;
    
    int argnum = lua_gettop(L);
    int argindex = 0;
    lua_Integer bytereq = 0;
    apr_size_t reqnow = 0;
    apr_size_t recvd = 0;
    const char *format = NULL;
    char bigbuff[LUA_SOCK_CHUNK_SIZE];
    char linebuff;
    
    luaL_Buffer b;
    
    if(argnum == 1) {
        /*If we didn't receive a format string, assume *l*/
        lua_pushstring(L, "*l");
        ++argnum;
    }
    
    for(argindex=2; argindex<=argnum; argindex++) {
        luaL_buffinit(L, &b);
        if(lua_type(L, argindex) == LUA_TNUMBER) {
            bytereq = luaL_checkinteger(L, argindex);
            if(bytereq <= 0) {
                lua_pushnil(L);
            } else {
                while(bytereq > 0) {
                    if(bytereq < LUA_SOCK_CHUNK_SIZE) {
                        reqnow = bytereq;
                    } else {
                        reqnow = LUA_SOCK_CHUNK_SIZE;
                    }
                    
                    rv = apr_socket_recv(sock->socket, bigbuff, &reqnow);
                    if(!reqnow) {
                        /*If we reach the end of the socket, return what we've got*/
                        break;
                    } else {
                        luaL_addlstring(&b, bigbuff, reqnow);
                    }
                    
                    bytereq -= reqnow;
                }
                
                luaL_pushresult(&b);
            }
        } else {
            format = luaL_checkstring(L, argindex);
            if(apr_strnatcmp(format, "*a") == 0) {
                /*read everything left on the socket*/
                while(rv != APR_EOF) {
                    recvd = LUA_SOCK_CHUNK_SIZE;
                    rv = apr_socket_recv(sock->socket, bigbuff, &recvd);
                    luaL_addlstring(&b, bigbuff, recvd);
                }
                luaL_pushresult(&b);
            } else {
                /*read the next line*/
                while(1) {
                    reqnow = 1;
                    rv = apr_socket_recv(sock->socket, &linebuff, &reqnow);
                    if(rv == APR_EOF || linebuff == '\n') {
                        break;
                    }
                    if(linebuff == '\r') continue;
                    
                    luaL_addchar(&b, linebuff);
                }
                
                luaL_pushresult(&b);
            }
        }
    }
    return argnum - 1;
}

/** Create a new raw socket and return it as a lua_apr_socket struct
 * \luareturn the raw apr.net.socket struct
 */
static int lua_apr_socket_new(lua_State *L) {
    apr_status_t rv;
    apr_pool_t *pool = NULL;
    
    rv = apr_pool_create(&pool, NULL);
    if(rv != APR_SUCCESS) {
        return luaL_error(L, "An error occurred initializing memory");
    }
    
    apr_socket_t *socket = NULL;
    rv = apr_socket_create(&socket, APR_INET, SOCK_STREAM, APR_PROTO_TCP, pool);
    if(rv != APR_SUCCESS) {
        return luaL_error(L, "An error occurred creating a new socket");
    }
    
    lua_apr_socket *sock = (lua_apr_socket *) lua_newuserdata(L, sizeof(lua_apr_socket));
    sock->socket = socket;
    sock->addr = NULL;
    sock->pool = pool;
    
    luaL_getmetatable(L, "apr.net.socket");
    lua_setmetatable(L, -2);
    return 1;
}

/** Bind a socket to a local address and port
 * \luaparam sock the apr.net.socket struct to bind
 * \luaparam port the port to bind to
 * \luaparam host the local interface to bind to - setting this to nil will bind to all interfaces
 */
static int lua_apr_socket_bind(lua_State *L) {
    lua_apr_socket *sock = (lua_apr_socket *) luaL_checkudata(L, 1, "apr.net.socket");
    apr_port_t port = (apr_port_t) luaL_checkint(L, 2);
    const char *host = lua_tostring(L, 3);
    apr_sockaddr_t *sa = NULL;
    apr_status_t rv;
    char errbuff[256];
    
    rv = apr_sockaddr_info_get(&sa, host, APR_UNSPEC, port, APR_IPV4_ADDR_OK, sock->pool);
    if(rv != APR_SUCCESS) {
        return luaL_error(L, "An error occurred getting the socket info: %s", apr_strerror(rv, errbuff, 256));
    }
    
    sock->addr = sa;
    
    rv = apr_socket_bind(sock->socket, sa);
    if(rv != APR_SUCCESS) {
        if(host == NULL) {
            host = "ANY";
        }
        
        return luaL_error(L, "Binding to %s:%d failed: %s", host, port, apr_strerror(rv, errbuff, 256));
    }
    
    return 0;
}

/** Listen to a bound socket for connections
 * \luaparam sock the apr.net.socket struct to bind
 * \luaparam backlog the number of backlogged requests to queue
 */
static int lua_apr_socket_listen(lua_State *L) {
    lua_apr_socket *sock = (lua_apr_socket *) luaL_checkudata(L, 1, "apr.net.socket");
    apr_int32_t backlog = (apr_int32_t) luaL_checkint(L, 2);
    apr_status_t rv;
    
    rv = apr_socket_listen(sock->socket, backlog);
    if(rv != APR_SUCCESS) {
        return luaL_error(L, "An error listening to the socket");
    }
    
    return 0;
}

/** Create a new client socket to use for the current request
 * \luaparam sock the server socket to accept from
 * \luareturn the new apr.net.socket struct for the request
 */
static int lua_apr_socket_accept(lua_State *L) {
    lua_apr_socket *sock = (lua_apr_socket *) luaL_checkudata(L, 1, "apr.net.socket");
    apr_status_t rv;
    apr_pool_t *pool = NULL;
    apr_socket_t *socket = NULL;
    
    rv = apr_pool_create(&pool, NULL);
    if(rv != APR_SUCCESS) {
        return luaL_error(L, "An error occurred initializing memory");
    }
    
    rv = apr_socket_accept(&socket, sock->socket, pool);
    if(rv != APR_SUCCESS) {
        return luaL_error(L, "An error occurred accepting client request");
    }
    
    lua_apr_socket *client_sock = (lua_apr_socket *) lua_newuserdata(L, sizeof(lua_apr_socket));
    client_sock->socket = socket;
    client_sock->addr = NULL;
    client_sock->pool = pool;
    
    luaL_getmetatable(L, "apr.net.socket");
    lua_setmetatable(L, -2);
    return 1;
}

static const struct luaL_reg lua_apr_rawsocket [] = {
    {"connect", lua_apr_socket_connect},
    {"close", lua_apr_socket_close},
    {"dnsresolve", lua_apr_socket_dnsresolve},
    {"send", lua_apr_socket_send},
    {"receive", lua_apr_socket_recv},
    {"newsocket", lua_apr_socket_new},
    {"bind", lua_apr_socket_bind},
    {"listen", lua_apr_socket_listen},
    {"accept", lua_apr_socket_accept},
    {NULL, NULL}
};

/** Set up the metatables and initialize the apr_dbd libraries
 */
int luaopen_apr_net_rawsocket(lua_State *L) {
    apr_status_t rv;
    luaL_register(L, "apr.net.rawsocket", lua_apr_rawsocket);
    
    rv = apr_initialize();
    if(rv != APR_SUCCESS) {
        return luaL_error(L, "Error initializing APR");
    }
    
    luaL_newmetatable(L, "apr.net.socket");
    lua_pushstring(L, "__gc");
    lua_pushcfunction(L, lua_apr_socket_close);
    lua_settable(L, -3);
    lua_pop(L, 1);
    
    return 1;
}
